import networkx as nx
import matplotlib.pyplot as plt
from graph import SpaceType
import numpy as np
from collections import defaultdict

class WisdomVisualizer:
    def __init__(self, wisdom_graph):
        self.wisdom = wisdom_graph
        self.colors = {
            SpaceType.PERCEPTION: '#ff9999',  # Light red
            SpaceType.PLANNING: '#99ff99',    # Light green
            SpaceType.REASONING: '#9999ff',   # Light blue
            SpaceType.ACTION: '#ffff99'       # Light yellow
        }

    def plot_wisdom_graph(self):
        """Visualize the entire wisdom graph with four distinct circular spaces"""
        plt.figure(figsize=(15, 15))
        
        # Define centers for each space region with more spread
        space_centers = {
            SpaceType.PERCEPTION: (0, 1.2),    # Top (moved up)
            SpaceType.PLANNING: (1.2, 0),      # Right (moved right)
            SpaceType.REASONING: (0, -1.2),    # Bottom (moved down)
            SpaceType.ACTION: (-1.2, 0)        # Left (moved left)
        }
        
        # Calculate positions within each space
        pos = {}
        for space_type in SpaceType:
            graph = getattr(self.wisdom, f"{space_type.value}_graph")
            center = space_centers[space_type]
            
            # Position nodes in a partial circle within their space (avoid center)
            node_count = len(graph.nodes())
            if node_count > 0:
                for i, node in enumerate(graph.nodes()):
                    # Use only part of the circle (180 degrees) facing outward
                    base_angle = {
                        SpaceType.PERCEPTION: np.pi/2,    # Top: point upward
                        SpaceType.PLANNING: 0,            # Right: point rightward
                        SpaceType.REASONING: -np.pi/2,    # Bottom: point downward
                        SpaceType.ACTION: np.pi           # Left: point leftward
                    }[space_type]
                    
                    angle = base_angle + (np.pi * (i - (node_count-1)/2) / (node_count + 1))
                    radius = 0.4  # Increased radius for more spread
                    x = center[0] + radius * np.cos(angle)
                    y = center[1] + radius * np.sin(angle)
                    pos[node] = (x, y)
        
        # Draw space regions (big circles)
        circle_radius = 0.6  # Slightly larger circles
        for space_type, center in space_centers.items():
            circle = plt.Circle(center, circle_radius, 
                              color=self.colors[space_type],
                              alpha=0.15)  # More transparent
            plt.gca().add_patch(circle)
            
            # Add space labels with larger font
            plt.text(center[0], center[1], space_type.value,
                    horizontalalignment='center',
                    verticalalignment='center',
                    fontsize=14,
                    alpha=0.7)
        
        # Draw curved edges showing the cognitive loop
        nx.draw_networkx_edges(self.wisdom.cross_space_edges, pos,
                              edge_color='gray',
                              arrows=True,
                              arrowsize=25,  # Larger arrows
                              connectionstyle="arc3,rad=0.2",  # Less curved
                              alpha=0.4,     # More transparent
                              width=2)       # Thicker lines
        
        # Draw nodes
        for space_type in SpaceType:
            graph = getattr(self.wisdom, f"{space_type.value}_graph")
            nx.draw_networkx_nodes(graph, pos,
                                 node_color=self.colors[space_type],
                                 node_size=1000,  # Larger nodes
                                 alpha=0.9,
                                 edgecolors='white',  # White border
                                 linewidths=2)
        
        # Add node labels with better positioning
        labels = {}
        for space_type in SpaceType:
            graph = getattr(self.wisdom, f"{space_type.value}_graph")
            for node in graph.nodes():
                cycle_num = node.split('_')[1]
                labels[node] = f"C{cycle_num}"
        
        nx.draw_networkx_labels(self.wisdom.cross_space_edges, pos, labels,
                              font_size=10,  # Larger font
                              font_weight='bold',
                              font_color='black')
        
        plt.title("Wisdom Graph Visualization\nCognitive Loop Across Four Spaces",
                  pad=20, fontsize=16, fontweight='bold')
        plt.gca().set_aspect('equal')
        plt.axis('off')
        
        # Add more margin around the plot
        plt.margins(0.2)
        plt.tight_layout()
        plt.show()

    def plot_clusters(self):
        """Visualize perception clusters"""
        if not self.wisdom.cluster_importance:
            return
            
        plt.figure(figsize=(10, 6))
        clusters = list(self.wisdom.cluster_importance.keys())
        importance = list(self.wisdom.cluster_importance.values())
        
        plt.bar(clusters, importance, color='skyblue')
        plt.title('Perception Clusters Importance')
        plt.xlabel('Cluster ID')
        plt.ylabel('Importance Score')
        plt.show()

class CognitiveVisualizer:
    def __init__(self):
        self.history = defaultdict(list)
        self.cycle_count = 0
        
        # Setup the plot
        plt.ion()  # Enable interactive mode
        self.fig, self.axes = plt.subplots(2, 2, figsize=(12, 8))
        self.fig.suptitle('Cognitive Cycle Metrics')
        
    def record_cycle(self, position, energy, knowledge, confidence):
        """Record metrics for each cognitive cycle"""
        self.cycle_count += 1
        self.history['position'].append(position)
        self.history['energy'].append(energy)
        self.history['knowledge'].append(knowledge)
        self.history['confidence'].append(confidence)
        
    def update_plot(self):
        """Update the live plot with current metrics"""
        self.axes[0, 0].clear()
        self.axes[0, 1].clear()
        self.axes[1, 0].clear()
        self.axes[1, 1].clear()
        
        cycles = range(1, self.cycle_count + 1)
        
        # Plot position components
        pos_x = [p[0] for p in self.history['position']]
        pos_y = [p[1] for p in self.history['position']]
        pos_z = [p[2] for p in self.history['position']]
        self.axes[0, 0].plot(cycles, pos_x, 'r-', label='X')
        self.axes[0, 0].plot(cycles, pos_y, 'g-', label='Y')
        self.axes[0, 0].plot(cycles, pos_z, 'b-', label='Z')
        self.axes[0, 0].set_title('Position Components')
        self.axes[0, 0].legend()
        
        # Plot energy
        self.axes[0, 1].plot(cycles, self.history['energy'], 'r-')
        self.axes[0, 1].set_title('Energy Level')
        
        # Plot knowledge
        self.axes[1, 0].plot(cycles, self.history['knowledge'], 'g-')
        self.axes[1, 0].set_title('Knowledge Level')
        
        # Plot confidence
        self.axes[1, 1].plot(cycles, self.history['confidence'], 'b-')
        self.axes[1, 1].set_title('Confidence Level')
        
        plt.tight_layout()
        plt.pause(0.1) 